////////////////////////////////////////////////////////////////
// multinomial expression divergence & gaussian
// (c) Yongjin Park, 2013
#include "med.hh"
#include <boost/math/special_functions.hpp>
#include <vector>

#ifndef MED_GAUSS_HH_
#define MED_GAUSS_HH_

struct
func_expr_pairwise_div_gauss_t : public func_expr_pairwise_div_t
{

public:

  typedef std::vector<double> vec_type;
  typedef func_expr_pairwise_div_t::data_t data_t;
  typedef func_expr_pairwise_div_t::data_set_t data_set_t;

private:

  // model for univariate Gaussian
  // x ~ N(mean,1/r)
  // mean ~ N(m0, 1/(s*r))
  // r ~ Gamma(c0, d0)

  // global hyper-parameters

  double m0;   // prior for mu ~ N(m0,s*r)
  double s;    // 
  double c0;   // prior for r ~ Gamma(c0,d0)
  double d0;   //


  // variational parameters for each type t

  vec_type mu;     // mean[t] ~ N(mu[t], gamma[t]^-1)
  vec_type gamma;  // 

  double c;        // r ~ Gamma(c, d)
  double d;        // 

  // expected sufficient statistics

  vec_type stat_nt;  // sum z_i       for each type t
  vec_type stat_St;  // sum z_i x_i   for each type t
  vec_type stat_S2t; // sum z_i x_i^2 for each type t

public:
  
  explicit func_expr_pairwise_div_gauss_t(const size_t p, double bmin = -5., double bmax = 5.) :
    func_expr_pairwise_div_t( p, bmin, bmax ), // pairwise divergence
    m0(0.), s(1.), c0(1.), d0(1.),             // hyper-parameters
    mu(p,m0), gamma(p,c0/d0), c(c0), d(d0),    // variational parameters
    stat_nt(p,0.), stat_St(p,0.), stat_S2t(p,0.)
    {}
  
  double update(const data_set_t& data_set,
		const double rate = 1.,
		const double size_factor = 1.)
    {
      double ret = func_expr_pairwise_div_t::update( data_set, rate, size_factor );

      // discount stat
      for(size_t t=0; t<num_vectors; ++t)
      {
	stat_nt[t] *= (1. - rate);
	stat_St[t] *= (1. - rate);
	stat_S2t[t] *= (1. - rate);
      }

      // update stat

      for(iterator_t it = begin(data_set); it != end(data_set); ++it)
      {
	const double z = it.get_prob();
	const data_t& data = *it;
	
	for(size_t t=0; t<num_vectors; ++t)
	{
	  const vec_type& X = data.vectors.at(t);
	  const double nt = X.size();

	  stat_nt[t] += rate * size_factor * nt * z;
	  
	  for(size_t a=0; a<X.size(); ++a)
	  {
	    const double x = X.at(a);
	    stat_St[t] += rate * size_factor * z * x;
	    stat_S2t[t] += rate * size_factor * z * x * x;
	  }
	}

      }

// #ifdef DEBUG
//       dump_vec( stat_St );
//       dump_vec( stat_nt );
// #endif

      // alternate update of N(mean|mu,gamma) and Gamma(r|c,d)
      // until convergence of RSS

      const size_t max_iter = 1000;
      const double tol = 1e-3;

      double RSS = 0.;

      double c_old = c;
      double d_old = d;

      for(size_t iter = 0; iter < max_iter; ++iter )
      {
	double RSS_old = RSS;

	// udpate RSS
	// RSS = sum_t stat_S2t - 2 * sum_t mu_t * ( m0 + stat_St )
	//       + sum_t (mu_t * mu_t + 1/E[gamma_t]) * (s + stat_nt)
	RSS = std::accumulate( stat_S2t.begin(), stat_S2t.end(), 0. );
	RSS += - 2. * m0 * std::accumulate( mu.begin(), mu.end(), 0. );
	RSS += - 2. * std::inner_product( mu.begin(), mu.end(), stat_St.begin(), 0. );
	for( size_t t=0; t<num_vectors; ++t)
	{
	  double mu2 = mu.at(t) * mu.at(t);
	  double var = 1./gamma.at(t);
	  RSS += (mu2 + var) * (s + stat_nt.at(t));
	}
      
	// update c and d
	// c = c0 + T/2 + stat_nt / 2
	// d = d0 + RSS / 2
	c = c0 + ((double) num_vectors) * 0.5 +
	  std::accumulate(stat_nt.begin(), stat_nt.end(), 0.) * 0.5;
	d = d0 + RSS * 0.5;

	// udpate mu and gamma
	// mu[t] = (stat_St + s * m0) / (stat_nt + s)
	// gamma[t] = E[r] * (stat_nt + s)
	//          = (c/d) * (stat_nt + s)
	for( size_t t=0; t<num_vectors; ++t )
	{
	  mu[t] = (stat_St.at(t) + s * m0) / (stat_nt.at(t) + s);
	  gamma[t] = c/d * (s + stat_nt.at(t));
	}

	if( std::sqrt( std::pow(RSS - RSS_old, 2.) ) < tol )
	{
#ifdef DEBUG
 	  std::cerr << "converged iter = " << iter
		    << " RSS = " << RSS << std::endl;
#endif
	  break;
	}
      }

// #ifdef DEBUG
//       std::cerr << "mu:"; dump_vec(mu);
//       std::cerr << "gamma:"; dump_vec(gamma);
// #endif

      ret += std::sqrt( std::pow(c-c_old, 2.) )
	+ std::sqrt( std::pow(d-d_old, 2.) );

      return ret;
    }

  double score(const data_t& data) const
    {
      double ret = func_expr_pairwise_div_t::score( data );
      ret += score_gauss( data );
      return ret;
    }

public:
  void set_m0( double _m0 ){ m0 = _m0; }
  // void set_scale( double _s ){ s = _s; }
  void set_c0( double _c0 ){ c0 = _c0; }
  void set_d0( double _d0 ){ d0 = _d0; }
  double get_m0() const { return m0; }
  // double get_scale() const { return s; }
  double get_c0() const { return c0; }
  double get_d0() const { return d0; }

  const vec_type& get_mu() const { return mu; }
  const vec_type& get_gamma() const { return gamma; }

  double get_prec() const { return ((c+c0)/(d+d0)); }

  void
  write_mu( std::ofstream& ofs )
    const
    {
      for(size_t k=0; k<num_vectors; ++k)
	ofs << k << "\t" << mu.at(k)
	    << "\t" << std::sqrt( 1. / gamma.at(k) )
	    << std::endl;
    }

private:
  
  double score_gauss( const data_t& data ) const
    {
      double ret = 0.;
      double D = 0.;
      double ntot = 0.;
      for(size_t t=0; t<num_vectors; ++t)
      {
	const vec_type& X = data.vectors.at(t);
	double nt = data.size();
	ntot += nt;
	double st = s + stat_nt.at(t);
	double mt = ( std::accumulate( X.begin(), X.end(), 0. ) + st * mu.at(t) );
	mt /= (nt + st);
	D += std::inner_product( X.begin(), X.end(), X.begin(), 0. );
	D += st * mu.at(t) * mu.at(t);
	D += - mt * mt * (nt + st);
      }

      ret += c * std::log( d );
      ret += - (c + 0.5 * ntot) * std::log(d + 0.5 * D);
      ret += boost::math::lgamma(ntot * 0.5 + c);
      ret += - boost::math::lgamma(c);
      
      // for(size_t t=0; t<num_vectors; ++t)
      // {
      // 	double nt = data.size();
      // 	double st = s + stat_nt.at(t);
      // 	ret += 0.5 * ( std::log(st) - std::log(nt + st) );	
      // }
      
      return ret;
    }

}; // end of med gauss

double empirical_bayes( func_expr_pairwise_div_gauss_t& obj,
			double rate = 1. )
{
  double ret = empirical_bayes_med( obj, rate );

  // too small number of observations ?
  // EB on Gaussian-Gamma may not work properly

  // update by pseudo-MLE
  // c0 = 1
  // d0 = 1 / E[r]  <=  E[r] = c0 / d0
  // where
  // E[r] = c/d
  // or
  // E[r] = ( sum_t gamma_t / s ) / T

  // double c0 = obj.get_c0();
  // double d0 = obj.get_d0();

  // typedef func_expr_pairwise_div_gauss_t::vec_type vec_type;
  // const vec_type& mu = obj.get_mu();
  // const vec_type& gamma = obj.get_gamma();

  // double T = mu.size();
  // double c0new = 1.;
  // double d0new = 1. / obj.get_prec();
    
  // c0 = (1.-rate) * c0 + rate * c0new;
  // d0 = (1.-rate) * d0 + rate * d0new;

  // obj.set_c0(c0);
  // obj.set_d0(d0);

  // ret += std::sqrt( std::pow(d0 - d0new, 2.) );

  return ret;
}

#endif
